from PIL import Image
from torchvision import transforms as T
import torchvision
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader, random_split
import pytorch_lightning as pl



class DataM(pl.LightningDataModule):
    def __init__(self, data_dir: str,
                 batch_size: int,
                 num_workers: int ,
                 resize: int):
        super().__init__()
        self.data_dir = data_dir
        self.batch_size = batch_size
        self.num_workers = num_workers
        self.resize = resize
        self.transform = T.Compose([
            T.ToTensor(),
            # T.RandomRotation(15),
            T.CenterCrop(self.resize),
            # T.Resize((int(self.resize * 1.25), int(self.resize * 1.25))),
            T.Normalize(mean=[0.485, 0.456, 0.406],
                        std=[0.229, 0.224, 0.225])
        ])
    def prepare_data(self):
        # download
        print("------------------prepare_data--------------")
        pass

    def setup(self, stage=None):
        dir_train = self.data_dir+"/train/"
        dir_test = self.data_dir+"/test/"
        if stage == 'fit' or stage is None:
            ds_full = torchvision.datasets.ImageFolder(root=dir_train,transform=self.transform)
            l = len(ds_full)
            l1 = int(0.1*l)
            self.ds_train, self.ds_val = random_split(ds_full, [l-l1, l1])
        if stage == 'test' or stage is None:
            self.ds_test = torchvision.datasets.ImageFolder(root=dir_test,transform=self.transform)

    def train_dataloader(self):
        return DataLoader(self.ds_train, batch_size=self.batch_size,
                          shuffle=True, num_workers=self.num_workers,
                          pin_memory=True)

    def val_dataloader(self):
        return DataLoader(self.ds_val, batch_size=self.batch_size,
                          shuffle=False, num_workers=self.num_workers,
                          pin_memory=True)

    def test_dataloader(self):
        return DataLoader(self.ds_test, batch_size=self.batch_size,
                          shuffle=False, num_workers=self.num_workers,
                          pin_memory=True)


if __name__ == '__main__':

    data_mnist = DataM("../../../../data/kaggle/Alzheimer's Dataset ( 4 class of Images)/Alzheimer_s Dataset/", 32, 0,224)
    data_mnist.setup()
    images, labels = next(iter(data_mnist.train_dataloader()))
    print(images.shape,labels)
    T.ToPILImage()(images[0]).show()





